package com.wuxuan.algorithm_homework.treeNode;

import org.junit.jupiter.api.Test;

import java.util.*;

/**
 * 二叉树 I
 * 1.前序遍历
 * 2.中序遍历
 * 3.后序遍历
 * 4.层序遍历
 * 5.Z字形遍历
 */
public class TreeNodeDemo {
    /**
     * 前序遍历 —— 递归
     * @param root
     */
    public void preOrderTraversal1(TreeNode root) {
        if (root == null) {
            return;
        }
        //也可以使用相应数据结构来存储，此处打印
        System.out.print(root.val + " ");
        preOrderTraversal1(root.left);
        preOrderTraversal1(root.right);

    }

    /**
     * 前序遍历 —— 迭代
     * @param root
     */
    public List<Integer> preOrderTraversal2(TreeNode root) {
        List<Integer> result = new ArrayList<>();
        if (root == null) {
            return result;
        }

        Deque<TreeNode> stack = new LinkedList<>();
        stack.offerLast(root);
        while (!stack.isEmpty()) {
            TreeNode cur = stack.pollLast();
            result.add(cur.val);

            if (cur.right != null) {
                stack.offerLast(cur.right);
            }

            if (cur.left != null) {
                stack.offerLast(cur.left);
            }
        }

        return result;

    }


    @Test
    public void testPreOrderTraversal() {
        //自创一个二叉树
        TreeNode root = new TreeNode(12);
        root.left = new TreeNode(5);
        root.right = new TreeNode(15);
        root.left.left = new TreeNode(3);
        root.left.right = new TreeNode(7);
        root.right.left = new TreeNode(13);
        root.right.right = new TreeNode(20);

        preOrderTraversal1(root);

        System.out.println();

        System.out.println(preOrderTraversal2(root));
    }

    /**
     * 中序遍历 —— 递归
     * @param root
     */
    public void inOrderTraversal1(TreeNode root) {
        if (root == null) {
            return;
        }
        inOrderTraversal1(root.left);
        System.out.print(root.val + " ");
        inOrderTraversal1(root.right);

    }

    /**
     * 中序遍历 —— 迭代
     * @param root
     */
    public List<Integer> inOrderTraversal2(TreeNode root) {
        List<Integer> result = new ArrayList<>();
        if (root == null) {
            return result;
        }

        Deque<TreeNode> stack = new LinkedList<>();
        TreeNode cur = root;
        while (cur != null || !stack.isEmpty()) {
            if (cur != null) {
                stack.offerLast(cur);
                cur = cur.left;
            } else {
                cur = stack.pollLast();
                result.add(cur.val);
                cur = cur.right;
            }
        }

        return result;

    }


    @Test
    public void testInOrderTraversal() {
        //自创一个二叉树
        TreeNode root = new TreeNode(12);
        root.left = new TreeNode(5);
        root.right = new TreeNode(15);
        root.left.left = new TreeNode(3);
        root.left.right = new TreeNode(7);
        root.right.left = new TreeNode(13);
        root.right.right = new TreeNode(20);

        inOrderTraversal1(root);

        System.out.println();

        System.out.println(inOrderTraversal2(root));
    }



    /**
     * 后序遍历 —— 递归
     * @param root
     */
    public void postOrderTraversal1(TreeNode root) {
        if (root == null) {
            return;
        }
        postOrderTraversal1(root.left);
        postOrderTraversal1(root.right);
        System.out.print(root.val + " ");

    }

    /**
     * 后序遍历 —— 迭代
     * @param root
     */
    public List<Integer> postOrderTraversal2(TreeNode root) {
        List<Integer> result = new ArrayList<>();
        if (root == null) {
            return result;
        }

        Deque<TreeNode> stack = new LinkedList<>();
        stack.offerLast(root);
        TreeNode cur = root;
        TreeNode pre = null;
        while (!stack.isEmpty()) {
            cur = stack.peekLast();
            if ((cur.left == null && cur.right == null) || (pre != null && pre == cur.left) || (pre != null && pre == cur.right)) {
                result.add(cur.val);
                stack.pollLast();
                pre = cur;
            } else {
                if (cur.right != null) {
                    stack.offerLast(cur.right);
                }

                if (cur.left != null) {
                    stack.offerLast(cur.left);
                }
            }
        }
        return result;
    }

    @Test
    public void testPostOrderTraversal() {
        //自创一个二叉树
        TreeNode root = new TreeNode(12);
        root.left = new TreeNode(5);
        root.right = new TreeNode(15);
        root.left.left = new TreeNode(3);
        root.left.right = new TreeNode(7);
        root.right.left = new TreeNode(13);
        root.right.right = new TreeNode(20);

        postOrderTraversal1(root);

        System.out.println();

        System.out.println(postOrderTraversal2(root));
    }


    /**
     * 层序遍历
     *
     * 树的层序遍历是一种对树进行遍历的方式，它按照树的层次依次访问每个节点
     * @param root
     * @return
     */
    public List<List<Integer>> levelOrder(TreeNode root) {
        List<List<Integer>> result = new ArrayList<>();
        if (root == null) {
            return result;
        }

        Queue<TreeNode> queue = new LinkedList<>();
        queue.offer(root);
        while (!queue.isEmpty()) {
            int size = queue.size();
            List<Integer> curResult = new ArrayList<>();
            while (size-- > 0) {
                TreeNode cur = queue.poll();
                curResult.add(cur.val);
                if (cur.left != null) {
                    queue.offer(cur.left);
                }
                if (cur.right != null) {
                    queue.offer(cur.right);
                }
            }
            result.add(curResult);
        }
        return result;
    }

    @Test
    public void testLevelOrder() {
        //自创一个二叉树
        TreeNode root = new TreeNode(12);
        root.left = new TreeNode(5);
        root.right = new TreeNode(15);
        root.left.left = new TreeNode(3);
        root.left.right = new TreeNode(7);
        root.right.left = new TreeNode(13);
        root.right.right = new TreeNode(20);

        System.out.println(levelOrder(root));
    }

    /**
     * Z字形层序遍历
     * @param root
     * @return
     */
    public List<List<Integer>> snakeTraversal(TreeNode root) {
        List<List<Integer>> result = new ArrayList<>();
        if (root == null) {
            return result;
        }
        Deque<TreeNode> deque = new LinkedList<>();
        deque.offerLast(root);
        int level = 0;
        while (!deque.isEmpty()) {
            int size = deque.size();
            level++;
            List<Integer> curResult = new ArrayList<>();
            while (size-- > 0) {
                TreeNode cur;
                if (level % 2 == 0) {
                    cur = deque.pollFirst();
                    if (cur.left != null) {
                        deque.offerLast(cur.left);
                    }
                    if (cur.right != null) {
                        deque.offerLast(cur.right);
                    }
                } else {
                    cur = deque.pollLast();
                    if (cur.right != null) {
                        deque.offerFirst(cur.right);
                    }
                    if (cur.left != null) {
                        deque.offerFirst(cur.left);
                    }
                }
                curResult.add(cur.val);
            }
            result.add(curResult);
        }
        return result;
    }

    @Test
    public void testSnakeTraversal() {
        //自创一个二叉树
        TreeNode root = new TreeNode(12);
        root.left = new TreeNode(5);
        root.right = new TreeNode(15);
        root.left.left = new TreeNode(3);
        root.left.right = new TreeNode(7);
        root.right.left = new TreeNode(13);
        root.right.right = new TreeNode(20);
        root.left.left.left = new TreeNode(0);
        root.left.right.right = new TreeNode(2);
        root.right.right.left = new TreeNode(1);

        System.out.println(snakeTraversal(root));
    }
}
